package com.zjj.lbw.ai.customer;

import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.net.URISyntaxException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.LocalDate;
import java.util.ArrayList;
import java.util.List;

/**
 * 智能客服系统，基于 AiService
 */
public class CustomerServiceAgent2 {
    public static void main(String[] args) {
        // 向量模型
        OpenAiEmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl("http://localhost:11434/v1")
                .apiKey("sk-a767d04a7e6f480bbd594c164c177775")
                .modelName("qwen2.5:14b")
                .build();

        // 向量存储数据库
        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        // 加载文本，并存储向量
        handleDocument(embeddingModel, embeddingStore);

        // 向量检索器
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(3) // 最相似的3个结果
                .minScore(0.7) // 只找相似度在0.8以上的内容
                .build();

        // 提示词注入器，用来将向量查询结果 封装为提示词，注入大模型
        DefaultContentInjector injector = new DefaultContentInjector();

        RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                .contentRetriever(contentRetriever)
                .contentInjector(injector)
                .build();
        // 大模型
//        ChatLanguageModel model = OpenAiChatModel.builder().modelName("gpt-4o-mini").apiKey("demo").build();

        ChatLanguageModel model = OpenAiChatModel.builder()
                .baseUrl("http://localhost:11434/v1")
                .modelName("qwen2.5:14b")
                .apiKey("sk-a767d04a7e6f480bbd594c164c177775")
                .build();

        AiCustomer aiCustomer = AiServices.builder(AiCustomer.class)
                .chatLanguageModel(model)
//                .retrievalAugmentor(retrievalAugmentor)
                .contentRetriever(contentRetriever)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .tools(new DateCalculator())
                .build();

        String answer = aiCustomer.call("今天的余额提现，最晚哪天能到账？给我具体的日期");
        System.out.println(answer);
    }

    private static void handleDocument(EmbeddingModel embeddingModel, EmbeddingStore embeddingStore) {
        // 加载并解析文件
        Document document;
        try {
            Path documentPath = Paths.get(CustomerServiceAgent.class.getClassLoader()
                    .getResource("meituan-qa.txt").toURI());
            DocumentParser documentParser = new TextDocumentParser();
            document = FileSystemDocumentLoader.loadDocument(documentPath, documentParser);

            // 切分文件
            DocumentSplitter splitter = new CustomerServiceAgent.CustomerServiceDocumentSplitter();
            List<TextSegment> segments = splitter.split(document);

            // 文本向量化 以及存储向量
            List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
            embeddingStore.addAll(embeddings, segments);
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
    }

    interface AiCustomer {
        String call(String query);
    }

    static class DateCalculator {

        @Tool("计算指定天数后的具体日期")
        String date(Integer days) {
            return LocalDate.now().plusDays(days).toString();
        }
    }

    static class CustomerServiceDocumentSplitter implements DocumentSplitter {

        @Override
        public List<TextSegment> split(Document document) {

            List<TextSegment> segments = new ArrayList<>();

            String[] parts = split(document.text());
            for (String part : parts) {
                segments.add(TextSegment.from(part));
            }

            return segments;
        }

        public String[] split(String text) {
            return text.split("\\s*\\R\\s*\\R\\s*");
        }

    }
}
